import pathlib

import scipy
import torch
import torch.utils.data


class CylinderNektarWakeDataset(torch.utils.data.Dataset):
    def __init__(self,
                 path: pathlib.Path = pathlib.Path("./data/cylinder_nektar_wake.mat")):
        x, y, t, u, v, p = CylinderNektarWakeDataset._load_mat(path)
        self._data = t, x, y, torch.stack([u, v, p], dim=-1)

    def __getitem__(self, item):
        t, x, y, f = self._data
        return t[item], x[item], y[item], f[item]

    def __len__(self):
        return self._data[0].shape[0]

    @staticmethod
    def _load_mat(path: pathlib.Path):
        # Modified from
        # https://github.com/maziarraissi/PINNs/blob/0542794b0a91b9e8764a38f5fc9cd9647a3929ba/main/continuous_time_identification%20(Navier-Stokes)/NavierStokes.py#L215
        data = scipy.io.loadmat(str(path))

        uv = torch.tensor(data["U_star"], dtype=torch.float)
        p = torch.tensor(data["p_star"], dtype=torch.float)
        t = torch.tensor(data["t"], dtype=torch.float)
        xy = torch.tensor(data["X_star"], dtype=torch.float)

        n_x = xy.shape[0]
        n_t = t.shape[0]

        x = torch.tile(xy[:, 0:1], (1, n_t))
        y = torch.tile(xy[:, 1:2], (1, n_t))
        t = torch.tile(t, (1, n_x)).T

        u = uv[:, 0, :]
        v = uv[:, 1, :]

        x = torch.flatten(x)
        y = torch.flatten(y)
        t = torch.flatten(t)

        u = torch.flatten(u)
        v = torch.flatten(v)
        p = torch.flatten(p)

        return x, y, t, u, v, p
